import polars as pl


from pathlib import Path


[f for f in Path("stock_trades/").glob("*-湘财.xls")]


def read_df_湘财(f: str | Path) -> pl.DataFrame:
    df = pl.read_csv(
        f,
        encoding="gb18030",
        separator="\t",
        infer_schema=False,
    )
    df = df.with_columns(
        pl.selectors.all().str.strip_prefix("=").str.strip_chars('"'),
    ).with_columns(
        pl.col("发生日期").str.to_date("%Y%m%d"),
        pl.col("成交时间").str.to_time(),
        pl.col(
            "成交数量",
            "成交价格",
            "成交金额",
            "发生金额",
            "手续费",
            "印花税",
            "过户费",
            "其他费",
        ).cast(pl.Float64),
    )
    df = df.filter(
        pl.col("业务名称").is_in(["证券买入", "证券卖出"]),
    )

    return df


d1 = pl.concat(df)


d1


df = [read_df_湘财(f) for f in Path("stock_trades/").glob("*-湘财.xls")]


d1 = pl.concat(df)


d1


d1 = d1.with_columns(
    券商=pl.lit("湘财"),
)


df = pl.read_excel(
    "stock_trades/202305-海通普通.xlsx",
    schema_overrides={
        "成交日期": pl.String,
        "成交时间": pl.String,
    },
)
df.filter(pl.col("成交时间") != "").filter(pl.col("操作").is_in(["买", "卖"])).filter(
    (~pl.col("证券代码").str.starts_with("204"))
    & (~pl.col("证券代码").str.starts_with("1318"))
).with_columns(
    pl.col("成交日期").str.to_date("%Y%m%d"),
    pl.col("成交时间").str.to_time("%H:%M:%S"),
)


def read_df_海通普通(f: str | Path) -> pl.DataFrame:
    df = pl.read_excel(
        f,
        schema_overrides={
            "成交日期": pl.String,
            "成交时间": pl.String,
            "成交数量": pl.Float64,
            "成交金额": pl.Float64,
            "印花税": pl.Float64,
            "其他费": pl.Float64,
        },
    )
    df = df.filter(
        (pl.col("成交时间") != "")
        & (pl.col("操作").is_in(["买", "卖"]))
        & (~pl.col("证券代码").str.starts_with("204"))
        & (~pl.col("证券代码").str.starts_with("1318"))
    ).with_columns(
        pl.col("成交日期").str.to_date("%Y%m%d"),
        pl.col("成交时间").str.to_time("%H:%M:%S"),
    )
    return df


df = [read_df_海通普通(p) for p in Path("stock_trades/").glob("*-海通普通.xlsx")]
df = pl.concat(df)
d2 = df.with_columns(券商=pl.lit("海通普通"))
d2.sort("成交日期", "成交时间")


df = [read_df_海通普通(p) for p in Path("stock_trades/").glob("*-海通两融.xlsx")]
df = pl.concat(df)
d3 = df.with_columns(券商=pl.lit("海通两融"))
d3.sort("成交日期", "成交时间")


d1 = d1.select(
    券商=pl.col("券商"),
    交易日期=pl.col("发生日期"),
    交易时间=pl.col("成交时间"),
    证券代码=pl.col("证券代码"),
    证券名称=pl.col("证券名称"),
    买卖标志=pl.col("业务名称").replace({"证券卖出": "卖出", "证券买入": "买入"}),
    成交价格=pl.col("成交价格"),
    成交数量=pl.col("成交数量").abs(),
    成交金额=pl.col("成交金额"),
    手续费=pl.col("手续费"),
    印花税=pl.col("印花税"),
    过户费=pl.col("过户费"),
    其他费=pl.col("其他费"),
    发生金额=pl.col("发生金额"),
)


d2


d2 = d2.select(
    券商=pl.col("券商"),
    交易日期=pl.col("成交日期"),
    交易时间=pl.col("成交时间"),
    证券代码=pl.col("证券代码"),
    证券名称=pl.col("证券名称"),
    买卖标志=pl.col("操作").replace({"卖": "卖出", "买": "买入"}),
    成交价格=pl.col("成交价格"),
    成交数量=pl.col("成交数量").abs(),
    成交金额=pl.col("成交金额"),
    手续费=pl.col("手续费"),
    印花税=pl.col("印花税"),
    过户费=pl.col("过户费"),
    其他费=pl.col("其他费"),
    发生金额=pl.col("发生金额"),
)


d3


d3 = d3.select(
    券商=pl.col("券商"),
    交易日期=pl.col("成交日期"),
    交易时间=pl.col("成交时间"),
    证券代码=pl.col("证券代码"),
    证券名称=pl.col("证券名称"),
    买卖标志=pl.col("操作").replace({"卖": "卖出", "买": "买入"}),
    成交价格=pl.col("成交价格"),
    成交数量=pl.col("成交数量").abs(),
    成交金额=pl.col("成交金额"),
    手续费=pl.col("手续费"),
    印花税=pl.col("印花税"),
    过户费=pl.col("过户费"),
    其他费=pl.col("其他费"),
    发生金额=pl.col("发生金额"),
)


df = pl.concat([d1, d2, d3])


df


# df.with_columns(
#     成交金额2=pl.col("成交金额") * pl.col("成交数量"),
# ).with_columns(
#     成交金额D=pl.col("成交金额") - pl.col("成交金额2"),
# ).with_columns(
#     发生金额D=pl.when(pl.col("买卖标志") == "买入")
#     .then(-pl.col("成交金额"))
#     .when(pl.col("买卖标志") == "卖出")
#     .then(pl.col("成交金额")),
# )


df.write_parquet("stock_tradesparquet")


df



